package cn.geminis.crypto.csp;

import cn.geminis.core.exception.CommonException;
import cn.geminis.crypto.core.util.ByteUtils;
import cn.geminis.crypto.core.x509.X509Certificate;
import org.bouncycastle.asn1.ASN1ObjectIdentifier;
import org.bouncycastle.asn1.x509.AlgorithmIdentifier;
import org.bouncycastle.cert.X509CertificateHolder;
import org.bouncycastle.cms.*;
import org.bouncycastle.cms.bc.BcKeyTransRecipientInfoGenerator;
import org.bouncycastle.crypto.AsymmetricBlockCipher;
import org.bouncycastle.crypto.BufferedBlockCipher;
import org.bouncycastle.crypto.io.CipherInputStream;
import org.bouncycastle.crypto.io.CipherOutputStream;
import org.bouncycastle.operator.GenericKey;
import org.bouncycastle.operator.InputDecryptor;
import org.bouncycastle.operator.OutputEncryptor;
import org.bouncycastle.operator.bc.BcAsymmetricKeyWrapper;

import java.io.Closeable;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;

/**
 * @author puddi
 */
public class Envelope implements Closeable {

    private AbstractBlockCipher symBlockCipher;
    private AbstractAsymmetricBlockCipher asymBlockCipher;
    private RandomGenerator randomGenerator;

    public Envelope(AbstractBlockCipher symBlockCipher, AbstractAsymmetricBlockCipher asymBlockCipher, RandomGenerator randomGenerator) {
        this.symBlockCipher = symBlockCipher;
        this.asymBlockCipher = asymBlockCipher;
        this.randomGenerator = randomGenerator;
    }

    @Override
    public void close() throws IOException {
        this.asymBlockCipher.close();
        this.symBlockCipher.close();
        this.randomGenerator.close();
    }

    public byte[] encrypt(byte[] data, byte[] cert) {
        var envelopeGenerator = new CMSEnvelopedDataGenerator();
        var x509Certificate = new X509Certificate(cert);
        var certHolder = new X509CertificateHolder(x509Certificate.getBcCertificate());
        try {
            envelopeGenerator.addRecipientInfoGenerator(new BcKeyTransRecipientInfoGenerator(
                    certHolder,
                    new BcAsymmetricKeyWrapper(
                            new AlgorithmIdentifier(new ASN1ObjectIdentifier(x509Certificate.getPublicKey().getAlgorithm())),
                            x509Certificate.getPublicKey().getKeyParameter()
                    ) {
                        @Override
                        protected AsymmetricBlockCipher createAsymmetricWrapper(ASN1ObjectIdentifier asn1ObjectIdentifier) {
                            return asymBlockCipher;
                        }
                    }) {
            });
        } catch (Exception e) {
            throw new CommonException("构建BcRSAKeyTransRecipientInfoGenerator错误", e);
        }

        data = ByteUtils.pkcs7Padding(data, symBlockCipher.getBlockSize());

        var msg = new CMSProcessableByteArray(data);
        try {
            var symKey = randomGenerator.random(symBlockCipher.getKeySize());
            symBlockCipher.init(true, symBlockCipher.createParams(symKey));

            var encryptor = new OutputEncryptor() {
                @Override
                public AlgorithmIdentifier getAlgorithmIdentifier() {
                    return new AlgorithmIdentifier(new ASN1ObjectIdentifier(symBlockCipher.getAlgorithmIdentifier()));
                }

                @Override
                public OutputStream getOutputStream(OutputStream outputStream) {
                    return new CipherOutputStream(outputStream, new BufferedBlockCipher(symBlockCipher));
                }

                @Override
                public GenericKey getKey() {
                    return new GenericKey(
                            new AlgorithmIdentifier(new ASN1ObjectIdentifier(symBlockCipher.getAlgorithmIdentifier())),
                            symKey);
                }
            };

            var envelope = envelopeGenerator.generate(msg, encryptor);
            return envelope.getEncoded();
        } catch (Exception e) {
            throw new CommonException("产生数字信封错误", e);
        }
    }

    public byte[] decrypt(byte[] data) {
        CMSEnvelopedData envelope;
        try {
            envelope = new CMSEnvelopedData(data);
        } catch (CMSException e) {
            throw new CommonException("解析数字信封错误", e);
        }

        var recipient = new KeyTransRecipient() {
            @Override
            public RecipientOperator getRecipientOperator(AlgorithmIdentifier algorithmIdentifier, AlgorithmIdentifier algorithmIdentifier1, byte[] bytes) {
                var decryptedSymKey = asymBlockCipher.decrypt(bytes);
                symBlockCipher.init(false, symBlockCipher.createParams(decryptedSymKey));
                return new RecipientOperator(new InputDecryptor() {
                    @Override
                    public AlgorithmIdentifier getAlgorithmIdentifier() {
                        return new AlgorithmIdentifier(new ASN1ObjectIdentifier(symBlockCipher.getAlgorithmIdentifier()));
                    }

                    @Override
                    public InputStream getInputStream(InputStream inputStream) {
                        return new CipherInputStream(inputStream, new BufferedBlockCipher(symBlockCipher));
                    }
                });
            }
        };

        var recipientInformation = (KeyTransRecipientInformation) envelope.getRecipientInfos().getRecipients().toArray()[0];
        try {
            var decryptedData = recipientInformation.getContent(recipient);
            return ByteUtils.pkcs7UnPadding(decryptedData);
        } catch (CMSException e) {
            throw new CommonException("解密数字信封错误", e);
        }
    }
}
